package websocket

import (
	"context"
	"encoding/json"
	"errors"
	"time"

	"github.com/gorilla/websocket"
)

const (
	// 默认时间设置
	defaultWriteWait  = 10 * time.Second
	defaultPongWait   = 60 * time.Second
	defaultPingPeriod = (defaultPongWait * 9) / 10
)

// NewConnection 创建新的 WebSocket 连接
func NewConnection(conn *websocket.Conn) *Connection {
	ctx, cancel := context.WithCancel(context.Background())
	return &Connection{
		conn:        conn,
		send:        make(chan []byte, 256),
		isClosed:    false,
		messageType: websocket.TextMessage,
		pingPeriod:  defaultPingPeriod,
		pongWait:    defaultPongWait,
		writeWait:   defaultWriteWait,
		ctx:         ctx,
		cancel:      cancel,
	}
}

// ReadPump 处理从 WebSocket 读取消息
func (c *Connection) ReadPump(messageHandler func([]byte) error) {
	defer func() {
		c.Close()
	}()

	c.conn.SetReadDeadline(time.Now().Add(c.pongWait))
	c.conn.SetPongHandler(func(string) error {
		c.conn.SetReadDeadline(time.Now().Add(c.pongWait))
		return nil
	})

	for {
		_, message, err := c.conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				// 非预期的关闭错误，可以记录日志
			}
			break
		}

		// 调用消息处理器
		if messageHandler != nil {
			if err := messageHandler(message); err != nil {
				// 处理消息时出错，可以选择关闭连接
				break
			}
		}
	}
}

// WritePump 处理向 WebSocket 写入消息
func (c *Connection) WritePump() {
	ticker := time.NewTicker(c.pingPeriod)
	defer func() {
		ticker.Stop()
		c.Close()
	}()

	for {
		select {
		case message, ok := <-c.send:
			c.conn.SetWriteDeadline(time.Now().Add(c.writeWait))
			if !ok {
				// Hub 关闭了通道
				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			w, err := c.conn.NextWriter(c.messageType)
			if err != nil {
				return
			}
			w.Write(message)

			// 添加队列中的其他消息
			n := len(c.send)
			for i := 0; i < n; i++ {
				w.Write([]byte{'\n'})
				w.Write(<-c.send)
			}

			if err := w.Close(); err != nil {
				return
			}

		case <-ticker.C:
			c.conn.SetWriteDeadline(time.Now().Add(c.writeWait))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}

		case <-c.ctx.Done():
			return
		}
	}
}

// Send 发送消息
func (c *Connection) Send(message []byte) error {
	c.mu.RLock()
	defer c.mu.RUnlock()

	if c.isClosed {
		return errors.New("connection is closed")
	}

	select {
	case c.send <- message:
		return nil
	case <-time.After(c.writeWait):
		return errors.New("send timeout")
	case <-c.ctx.Done():
		return errors.New("connection closed")
	}
}

// SendJSON 发送 JSON 消息
func (c *Connection) SendJSON(v interface{}) error {
	data, err := json.Marshal(v)
	if err != nil {
		return err
	}
	return c.Send(data)
}

// SendMessage 发送结构化消息
func (c *Connection) SendMessage(msgType string, payload interface{}) error {
	msg := Message{
		Type:    msgType,
		Payload: payload,
	}
	return c.SendJSON(msg)
}

// Close 关闭连接
func (c *Connection) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.isClosed {
		return nil
	}

	c.isClosed = true
	c.cancel()
	close(c.send)

	return c.conn.Close()
}

// IsClosed 检查连接是否已关闭
func (c *Connection) IsClosed() bool {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.isClosed
}

// SetPongWait 设置 Pong 等待时间
func (c *Connection) SetPongWait(d time.Duration) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.pongWait = d
	c.pingPeriod = (d * 9) / 10
}

// SetWriteWait 设置写入等待时间
func (c *Connection) SetWriteWait(d time.Duration) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.writeWait = d
}

// GetRemoteAddr 获取远程地址
func (c *Connection) GetRemoteAddr() string {
	if c.conn != nil {
		return c.conn.RemoteAddr().String()
	}
	return ""
}

